import numpy as np
import matplotlib.pyplot as plt
import math
import mpl_toolkits.axisartist as axisartist


def triangle_signal(f):
    if f < -50 or f > 50:
        return 0
    elif f < 0:
        return f / 50 + 1
    else:
        return 1 - f / 50

def ifft(t, delta):
    if t == 0:
        return delta
    else:
        return math.sin(math.pi * delta * (-t)) ** 2 / (math.pi ** 2 * delta * (-t) ** 2)

def sample_convert(t, Ts):
    amplitude = np.zeros_like(t)
    def sa(x):
        return math.sin(x) / x
    for i in range(t.size):
        for j in range(-10, 11):
            amplitude[i] += ifft(j * Ts, delta) * sa(math.pi / Ts * (t[i] - j * Ts))
    return amplitude
    

def set_axis(fig, number, title, x_range=(-10, 15), y_range=(-1, 1)):
    ax = axisartist.Subplot(fig, number)
    fig.add_axes(ax)
    ax.axis[:].set_visible(False)
    ax.axis["x"] = ax.new_floating_axis(0,0)
    ax.axis["x"].set_axisline_style("->", size = 1.0)
    ax.axis["y"] = ax.new_floating_axis(1,0)
    ax.axis["y"].set_axisline_style("->", size = 1.0)
    ax.axis["x"].set_axis_direction("bottom")
    ax.axis["y"].set_axis_direction("left")
    ax.set_xlim(x_range[0], x_range[1])
    ax.set_ylim(y_range[0], y_range[1])
    ax.set_title(title)


if __name__ == "__main__":

    fig = plt.figure(figsize=(8, 8))
    plt.rcParams['font.family'] = ['Arial Unicode MS']

    T = 100
    delta = 50
    t = np.arange(-0.1, 0.1, 0.001)

    time_signal = np.array([ifft(i, delta) for i in t])
    set_axis(fig, 311, "带限信号时间域图像", (-0.1, 0.1), (-2, 50))
    plt.plot(t, time_signal)

    Ts = 0.02
    time_signal = sample_convert(t, Ts)
    set_axis(fig, 312, "50Hz抽样后恢复的连续信号", (-0.1, 0.1), (-20, 50))
    plt.plot(t, time_signal)

    Ts = 1 / 150
    time_signal = sample_convert(t, Ts)
    set_axis(fig, 313, "150Hz抽样后恢复的连续信号", (-0.1, 0.1), (-2, 50))
    plt.plot(t, time_signal)

    plt.show()